"""
Database migration manager for CodeMCP storage layer.

Handles schema versioning, migration execution, and rollback operations
using Alembic for database schema management.
"""

import os
import asyncio
from pathlib import Path
from typing import Dict, Any, List, Optional
from datetime import datetime
import importlib.util

from alembic import command
from alembic.config import Config as AlembicConfig
from alembic.runtime.migration import MigrationContext
from alembic.operations import Operations
from alembic.script import ScriptDirectory
from sqlalchemy import create_engine, text
from sqlalchemy.ext.asyncio import create_async_engine

from ..database import DatabaseManager
from ...core.config import Config
from ...core.error_handler import CodeMCPError, log_info, log_warning, log_debug, handle_error


class MigrationError(CodeMCPError):
    """Migration-specific errors."""
    
    def __init__(self, message: str, operation: str = None, **kwargs):
        super().__init__(message, "MIGRATION_ERROR", **kwargs)
        if operation:
            self.details["operation"] = operation


class MigrationManager:
    """
    Database migration manager using Alembic.
    
    Handles database schema versioning, migration execution,
    and rollback operations for the CodeMCP storage layer.
    """
    
    def __init__(self, config: Config = None, db_manager: DatabaseManager = None):
        self.config = config or Config()
        self.db_manager = db_manager
        
        # Migration configuration
        self.migrations_dir = Path(__file__).parent
        self.alembic_cfg_path = self.migrations_dir / "alembic.ini"
        self.versions_dir = self.migrations_dir / "versions"
        
        # Alembic configuration
        self.alembic_cfg = None
        self.script_directory = None
        
        # Migration statistics
        self.stats = {
            "migrations_executed": 0,
            "rollbacks_executed": 0,
            "errors_encountered": 0,
            "last_migration": None
        }
        
        # Ensure directories exist
        self.versions_dir.mkdir(exist_ok=True)
        
        log_info("Migration manager initialized")
    
    async def initialize(self):
        """Initialize the migration manager."""
        try:
            # Create Alembic configuration if it doesn't exist
            if not self.alembic_cfg_path.exists():
                await self._create_alembic_config()
            
            # Load Alembic configuration
            self.alembic_cfg = AlembicConfig(str(self.alembic_cfg_path))
            
            # Set up script directory
            self.script_directory = ScriptDirectory.from_config(self.alembic_cfg)
            
            # Ensure database is initialized
            if self.db_manager and not self.db_manager.is_initialized:
                await self.db_manager.initialize()
            
            log_info("Migration manager initialized successfully")
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "migration_manager_init"})
            raise MigrationError(f"Failed to initialize migration manager: {error_response['message']}")
    
    async def _create_alembic_config(self):
        """Create Alembic configuration file."""
        alembic_ini_content = f"""
# Alembic configuration for CodeMCP

[alembic]
# Path to migration scripts
script_location = {self.migrations_dir}

# Template used to generate migration files
file_template = %%(year)d%%(month).2d%%(day).2d_%%(hour).2d%%(minute).2d_%%(rev)s_%%(slug)s

# Timezone to use when rendering the date within the migration file
timezone =

# Max length of characters to apply to the "slug" field
truncate_slug_length = 40

# Set to 'true' to run the environment during
# the 'revision' command, regardless of autogenerate
revision_environment = false

# Set to 'true' to allow .pyc and .pyo files without
# a source .py file to be detected as revisions in the
# versions/ directory
sourceless = false

# Version path separator; as mentioned above, this is the default
version_path_separator = os  # Use os.pathsep. Default: space

# The output encoding used when revision files
# are written from script.py.mako
output_encoding = utf-8

sqlalchemy.url = {self.db_manager.database_url if self.db_manager else 'sqlite:///codemcp.db'}

[post_write_hooks]
# Post-write hooks define scripts or Python functions that are run
# on newly generated revision scripts.

# Logging configuration
[loggers]
keys = root,sqlalchemy,alembic

[handlers]
keys = console

[formatters]
keys = generic

[logger_root]
level = WARN
handlers = console
qualname =

[logger_sqlalchemy]
level = WARN
handlers =
qualname = sqlalchemy.engine

[logger_alembic]
level = INFO
handlers =
qualname = alembic

[handler_console]
class = StreamHandler
args = (sys.stderr,)
level = NOTSET
formatter = generic

[formatter_generic]
format = %(levelname)-5.5s [%(name)s] %(message)s
datefmt = %H:%M:%S
"""
        
        # Write Alembic configuration
        with open(self.alembic_cfg_path, 'w') as f:
            f.write(alembic_ini_content.strip())
        
        # Create env.py file
        env_py_content = '''
"""
Alembic environment for CodeMCP migrations.
"""

import asyncio
from logging.config import fileConfig
from sqlalchemy import pool
from sqlalchemy.engine import Connection
from sqlalchemy.ext.asyncio import async_engine_from_config
from alembic import context

# Import your models here
from codemcp.storage.models import Base

# Alembic Config object
config = context.config

# Configure logging
if config.config_file_name is not None:
    fileConfig(config.config_file_name)

# Set target metadata
target_metadata = Base.metadata


def run_migrations_offline() -> None:
    """Run migrations in 'offline' mode."""
    url = config.get_main_option("sqlalchemy.url")
    context.configure(
        url=url,
        target_metadata=target_metadata,
        literal_binds=True,
        dialect_opts={"paramstyle": "named"},
    )

    with context.begin_transaction():
        context.run_migrations()


def do_run_migrations(connection: Connection) -> None:
    """Run migrations with database connection."""
    context.configure(connection=connection, target_metadata=target_metadata)

    with context.begin_transaction():
        context.run_migrations()


async def run_async_migrations() -> None:
    """Run migrations in async mode."""
    connectable = async_engine_from_config(
        config.get_section(config.config_ini_section, {}),
        prefix="sqlalchemy.",
        poolclass=pool.NullPool,
    )

    async with connectable.connect() as connection:
        await connection.run_sync(do_run_migrations)

    await connectable.dispose()


def run_migrations_online() -> None:
    """Run migrations in 'online' mode."""
    asyncio.run(run_async_migrations())


if context.is_offline_mode():
    run_migrations_offline()
else:
    run_migrations_online()
'''
        
        env_py_path = self.migrations_dir / "env.py"
        with open(env_py_path, 'w') as f:
            f.write(env_py_content.strip())
        
        # Create script.py.mako template
        script_mako_content = '''"""${message}

Revision ID: ${up_revision}
Revises: ${down_revision | comma,n}
Create Date: ${create_date}

"""
from alembic import op
import sqlalchemy as sa
${imports if imports else ""}

# Revision identifiers
revision = ${repr(up_revision)}
down_revision = ${repr(down_revision)}
branch_labels = ${repr(branch_labels)}
depends_on = ${repr(depends_on)}


def upgrade() -> None:
    """Upgrade database schema."""
    ${upgrades if upgrades else "pass"}


def downgrade() -> None:
    """Downgrade database schema."""
    ${downgrades if downgrades else "pass"}
'''
        
        script_mako_path = self.migrations_dir / "script.py.mako"
        with open(script_mako_path, 'w') as f:
            f.write(script_mako_content.strip())
        
        log_info("Created Alembic configuration files")
    
    async def get_current_revision(self) -> Optional[str]:
        """Get the current database revision."""
        try:
            if not self.db_manager:
                return None
            
            # Check if alembic_version table exists
            result = await self.db_manager.execute_query(
                "SELECT name FROM sqlite_master WHERE type='table' AND name='alembic_version'"
            )
            
            if not result.fetchone():
                return None
            
            # Get current revision
            result = await self.db_manager.execute_query(
                "SELECT version_num FROM alembic_version"
            )
            
            row = result.fetchone()
            return row[0] if row else None
            
        except Exception as e:
            log_warning(f"Failed to get current revision: {e}")
            return None
    
    async def get_head_revision(self) -> Optional[str]:
        """Get the head revision from migration scripts."""
        try:
            if not self.script_directory:
                return None
            
            return self.script_directory.get_current_head()
            
        except Exception as e:
            log_warning(f"Failed to get head revision: {e}")
            return None
    
    async def is_database_up_to_date(self) -> bool:
        """Check if database is up to date with latest migrations."""
        current = await self.get_current_revision()
        head = await self.get_head_revision()
        
        return current == head
    
    async def get_pending_migrations(self) -> List[str]:
        """Get list of pending migration revisions."""
        try:
            current = await self.get_current_revision()
            
            if not self.script_directory:
                return []
            
            if current is None:
                # No migrations applied yet, return all migrations
                revisions = []
                for revision in self.script_directory.walk_revisions():
                    revisions.append(revision.revision)
                return list(reversed(revisions))  # Return in execution order
            
            # Get revisions between current and head
            pending = []
            for revision in self.script_directory.walk_revisions():
                if revision.revision == current:
                    break
                pending.append(revision.revision)
            
            return pending
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "get_pending_migrations"})
            raise MigrationError(f"Failed to get pending migrations: {error_response['message']}")
    
    async def migrate_to_latest(self) -> bool:
        """Migrate database to the latest revision."""
        try:
            if not self.alembic_cfg:
                await self.initialize()
            
            # Run migrations
            command.upgrade(self.alembic_cfg, "head")
            
            self.stats["migrations_executed"] += 1
            self.stats["last_migration"] = datetime.utcnow().isoformat()
            
            log_info("Database migrated to latest revision")
            return True
            
        except Exception as e:
            self.stats["errors_encountered"] += 1
            error_response = handle_error(e, {"operation": "migrate_to_latest"})
            raise MigrationError(f"Failed to migrate to latest: {error_response['message']}")
    
    async def migrate_to_revision(self, revision: str) -> bool:
        """Migrate database to a specific revision."""
        try:
            if not self.alembic_cfg:
                await self.initialize()
            
            # Run migration to specific revision
            command.upgrade(self.alembic_cfg, revision)
            
            self.stats["migrations_executed"] += 1
            self.stats["last_migration"] = datetime.utcnow().isoformat()
            
            log_info(f"Database migrated to revision: {revision}")
            return True
            
        except Exception as e:
            self.stats["errors_encountered"] += 1
            error_response = handle_error(e, {
                "operation": "migrate_to_revision",
                "revision": revision
            })
            raise MigrationError(f"Failed to migrate to revision {revision}: {error_response['message']}")
    
    async def rollback_to_revision(self, revision: str) -> bool:
        """Rollback database to a specific revision."""
        try:
            if not self.alembic_cfg:
                await self.initialize()
            
            # Rollback to specific revision
            command.downgrade(self.alembic_cfg, revision)
            
            self.stats["rollbacks_executed"] += 1
            
            log_info(f"Database rolled back to revision: {revision}")
            return True
            
        except Exception as e:
            self.stats["errors_encountered"] += 1
            error_response = handle_error(e, {
                "operation": "rollback_to_revision",
                "revision": revision
            })
            raise MigrationError(f"Failed to rollback to revision {revision}: {error_response['message']}")
    
    async def create_migration(self, message: str, autogenerate: bool = True) -> Optional[str]:
        """Create a new migration script."""
        try:
            if not self.alembic_cfg:
                await self.initialize()
            
            # Create new migration
            if autogenerate:
                script = command.revision(
                    self.alembic_cfg,
                    message=message,
                    autogenerate=True
                )
            else:
                script = command.revision(
                    self.alembic_cfg,
                    message=message
                )
            
            log_info(f"Created migration: {script.revision} - {message}")
            return script.revision
            
        except Exception as e:
            self.stats["errors_encountered"] += 1
            error_response = handle_error(e, {
                "operation": "create_migration",
                "message": message
            })
            raise MigrationError(f"Failed to create migration: {error_response['message']}")
    
    async def get_migration_history(self) -> List[Dict[str, Any]]:
        """Get migration history."""
        try:
            if not self.script_directory:
                return []
            
            history = []
            for revision in self.script_directory.walk_revisions():
                history.append({
                    "revision": revision.revision,
                    "down_revision": revision.down_revision,
                    "branch_labels": revision.branch_labels,
                    "message": revision.doc,
                    "module_path": revision.path
                })
            
            return history
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "get_migration_history"})
            raise MigrationError(f"Failed to get migration history: {error_response['message']}")
    
    async def validate_database_schema(self) -> Dict[str, Any]:
        """Validate that database schema matches models."""
        try:
            # This would compare the current database schema with the SQLAlchemy models
            # For now, we'll do a basic check
            
            current_revision = await self.get_current_revision()
            head_revision = await self.get_head_revision()
            pending_migrations = await self.get_pending_migrations()
            
            is_valid = current_revision == head_revision
            
            return {
                "valid": is_valid,
                "current_revision": current_revision,
                "head_revision": head_revision,
                "pending_migrations": pending_migrations,
                "pending_count": len(pending_migrations)
            }
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "validate_database_schema"})
            raise MigrationError(f"Failed to validate database schema: {error_response['message']}")
    
    async def get_migration_stats(self) -> Dict[str, Any]:
        """Get migration statistics."""
        try:
            current_revision = await self.get_current_revision()
            head_revision = await self.get_head_revision()
            pending_migrations = await self.get_pending_migrations()
            migration_history = await self.get_migration_history()
            
            return {
                "migration_stats": self.stats.copy(),
                "current_revision": current_revision,
                "head_revision": head_revision,
                "pending_migrations": pending_migrations,
                "pending_count": len(pending_migrations),
                "total_migrations": len(migration_history),
                "is_up_to_date": await self.is_database_up_to_date()
            }
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "get_migration_stats"})
            return {
                "migration_stats": self.stats.copy(),
                "error": error_response["message"]
            }